// Copyright FYT 邹承甫 2022

#include "post_processor/tracker.hpp"

#include <bits/stdint-uintn.h>
#include <fmt/color.h>
#include <fmt/core.h>

#include "myutils/utils.hpp"

namespace post_processor {

Tracker::Tracker(const std::vector<double> &q_vec,
                 const std::vector<double> &p_vec,
                 const std::vector<double> &f_vec,
                 const std::vector<double> &h_vec,
                 const std::vector<double> &r_vec,
                 const double &max_match_distance,
                 const double &max_spinning_angle, const int &max_queue_length,
                 const bool &is_debug)
    : debug_(is_debug),
      max_match_thresh_(max_match_distance),
      max_queue_size_(max_queue_length),
      max_angle_(max_spinning_angle) {
  last_time = ros::Time::now();
  Eigen::MatrixXd Q = Eigen::MatrixXd::Zero(7, 7);
  Eigen::MatrixXd P = Eigen::MatrixXd::Zero(7, 7);
  Eigen::MatrixXd F = Eigen::MatrixXd::Zero(7, 7);
  Eigen::MatrixXd H = Eigen::MatrixXd::Zero(7, 7);
  Eigen::MatrixXd R = Eigen::MatrixXd::Zero(7, 7);
  utils::toEigenMatrix(Q, q_vec);
  utils::toEigenMatrix(P, p_vec);
  utils::toEigenMatrix(F, f_vec);
  utils::toEigenMatrix(H, h_vec);
  utils::toEigenMatrix(R, r_vec);
  kf_ = std::make_unique<KalmanFilter>(F, P, Q, R, H);

  predict_time_ = 0.005;
  lost_cnt_ = 0;
  tracking_id_ = 0;
  last_state = Eigen::VectorXd::Zero(7);
  state_ = tracker::state::LOST;
}

void Tracker::initialize(const std::vector<Armor> &targets) {
  Armor selected_armor = targets[0];
  double min_distance = 10e5;
  for (auto &armor : targets) {
    double distance = armor.position.norm();
    if (distance < min_distance) {
      selected_armor = armor;
      min_distance = distance;
    }
  }
  history_queue_.clear();

  Eigen::VectorXd state = Eigen::VectorXd::Zero(7);
  state << selected_armor.position(0), selected_armor.position(1),
      selected_armor.position(2), 0, 0, 0, selected_armor.yaw;

  kf_->initialize(state);

  lost_cnt_ = 0;
  state_ = tracker::TRACKING;
  tracking_id_ = selected_armor.id;
  history_queue_.push_back(selected_armor);
  last_time = ros::Time::now();
  last_state = Eigen::VectorXd::Zero(7);
}

Eigen::VectorXd Tracker::update(const std::vector<Armor> &targets) {
  // state : {x, y, z, vx, vy, vz, angle(yaw)}
  //掉帧处理
  if (targets.empty()) {
    Eigen::VectorXd state = Eigen::VectorXd::Zero(7);
    if (this->state_ == tracker::state::TRACKING) {
      lost_cnt_++;
      if (lost_cnt_ > 10) {
        this->state_ = tracker::state::LOST;
      }
      state = kf_->predict(lost_cnt_ * predict_time_);
      state[3] /= lost_cnt_;
      state[4] /= lost_cnt_;
      state[5] /= lost_cnt_;
      return state;
    } else {
      return state;
    }
  }
  //追踪间隔太长
  double dt = (targets[0].stamp - last_time).toSec();
  last_time = targets[0].stamp;
  if (state_ == tracker::state::LOST || dt > 1.0) {
    initialize(targets);
    Eigen::VectorXd target_state = kf_->predict(0);
    return target_state;
  }

  // 选择与上一帧跟踪的装甲板最近的作为目标
  bool matched = false;
  double min_diff = 10e9;
  Armor matched_armor = targets[0];
  Eigen::VectorXd predict_pos = kf_->predict(dt);
  for (auto &target : targets) {
    double diff = (predict_pos.head(3) - target.position).norm() +
                  fabs(predict_pos[6] - target.yaw);
    if (diff < min_diff) {
      matched_armor = target;
      min_diff = diff;
    }
  }

  // 考虑倾斜装甲板在PnP解算时误差较大，当装甲板倾斜且id没问题时，放宽匹配条件
  double match_dist_thresh = max_match_thresh_;
  if (matched_armor.id == tracking_id_ &&
      (matched_armor.yaw > max_angle_ || lost_cnt_ != 0)) {
    match_dist_thresh *= 2.0;
  }
  if (min_diff < match_dist_thresh) {
    matched = true;
    lost_cnt_ = 0;
  }

  if (matched) {
    // 计算速度
    double total_time =
        (matched_armor.stamp - history_queue_.front().stamp).toSec();
    Eigen::Vector3d velocity =
        (matched_armor.position - history_queue_.front().position) / total_time;
    velocity = history_queue_.size() > 1 ? velocity : 0.5 * velocity;

    Armor final_target = matched_armor;
    // 当目前跟踪的装甲板过于倾斜时，启动反小陀螺
    bool is_spinning = false;
    if (fabs(matched_armor.yaw) > max_angle_ && history_queue_.size() > 1) {
      is_spinning = antiSpinning(targets, matched_armor, final_target);
    }

    if (is_spinning) {
      if (debug_) {
        fmt::print(
            fg(fmt::color::purple) | fmt::emphasis::bold |
                fmt::emphasis::italic,
            "Spinning detected! new target is [{:.1f}, {:.1f}, {:.1f}]\n",
            final_target.position[0], final_target.position[1],
            final_target.position[2]);
      }
      history_queue_.clear();

      Eigen::VectorXd measurement = Eigen::VectorXd::Zero(7);
      //小陀螺时保持当前速度
      measurement << final_target.position[0], final_target.position[1],
          final_target.position[2], velocity[0], velocity[1], velocity[2],
          final_target.yaw;
      kf_->initialize(measurement);
      Eigen::VectorXd state = kf_->predict(0);

      history_queue_.emplace_back(
          state.head(3), final_target.stamp, final_target.id, state[6]);
      last_state = state;
      tracking_id_ = final_target.id;
      return state;
    } else {
      Eigen::VectorXd measurement = Eigen::VectorXd::Zero(7);
      measurement << final_target.position[0], final_target.position[1],
          final_target.position[2], velocity[0], velocity[1], velocity[2],
          final_target.yaw;
      Eigen::VectorXd state = kf_->update(measurement, dt);

      history_queue_.emplace_back(
          state.head(3), final_target.stamp, final_target.id, state(6));
      // 队满或者以当前队头计算出来的速度很离谱(卡尔曼滤波器的卡方检验不通过)时队头出队
      if (static_cast<int>(history_queue_.size()) > max_queue_size_ ||
          state.block(3, 0, 3, 1).norm() == 0) {
        history_queue_.pop_front();
      }
      last_state = state;
      tracking_id_ = final_target.id;
      state = kf_->predict(predict_time_);
      return state;
    }
  } else {
    // 没有匹配的装甲板时
    lost_cnt_++;
    //连续丢失5帧重新选择跟踪目标
    if (lost_cnt_ > 5) {
      initialize(targets);
      Eigen::VectorXd state = kf_->predict(0);
      return state;
    } else {
      Eigen::VectorXd state = kf_->predict(lost_cnt_ * dt + predict_time_);
      state[3] /= lost_cnt_;
      state[4] /= lost_cnt_;
      state[5] /= lost_cnt_;
      return state;
    }
  }
}

bool Tracker::antiSpinning(const std::vector<Armor> &armors,
                           const Armor &selected, Armor &new_target) {
  // 找有没有相同id的装甲板
  std::vector<Armor> same_id_armors;
  for (const auto &armor : armors) {
    if (armor.id == selected.id && armor.position != selected.position) {
      same_id_armors.emplace_back(armor);
    }
  }
  if (same_id_armors.empty()) {
    return false;
  }
  // 在相同id的装甲板中挑选yaw最小的
  double min_yaw = 10e9;
  Armor matched_armor = selected;
  for (auto &armor : same_id_armors) {
    if (fabs(armor.yaw) < min_yaw) {
      min_yaw = fabs(armor.yaw);
      matched_armor = armor;
    }
  }
  if (fabs(matched_armor.yaw) < fabs(selected.yaw)) {
    new_target = matched_armor;
    return true;
  } else {
    return false;
  }
}

void Tracker::set_predict_time(const double &t) {
  predict_time_ = t;
}

}  // namespace post_processor